package com.binance.feign.ribbon.lb;

import com.binance.feign.ribbon.exception.RibbonResponseStatusCodeException;
import com.binance.feign.ribbon.policy.FeignRetryPolicy;
import com.binance.feign.ribbon.policy.RibbonLoadBalancedRetryPolicy;
import com.netflix.client.DefaultLoadBalancerRetryHandler;
import com.netflix.client.RequestSpecificRetryHandler;
import com.netflix.client.config.CommonClientConfigKey;
import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import feign.Request;
import feign.Response;
import lombok.extern.log4j.Log4j2;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancedBackOffPolicyFactory;
import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryContext;
import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryListenerFactory;
import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryPolicy;
import org.springframework.cloud.client.loadbalancer.RibbonRecoveryCallback;
import org.springframework.cloud.client.loadbalancer.ServiceInstanceChooser;
import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerClient;
import org.springframework.cloud.netflix.ribbon.ServerIntrospector;
import org.springframework.cloud.netflix.ribbon.SpringClientFactory;
import org.springframework.retry.RetryListener;
import org.springframework.retry.backoff.BackOffPolicy;
import org.springframework.retry.backoff.NoBackOffPolicy;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.StreamUtils;

import java.io.IOException;
import java.net.URI;

/**
 * Created by Fei.Huang on 2018/9/7.
 */
@Log4j2
public class RetryableFeignLoadBalancer extends FeignLoadBalancer implements ServiceInstanceChooser {

    private SpringClientFactory clientFactory;

    private final LoadBalancedBackOffPolicyFactory loadBalancedBackOffPolicyFactory;
    private final LoadBalancedRetryListenerFactory loadBalancedRetryListenerFactory;

    public RetryableFeignLoadBalancer(ILoadBalancer lb, IClientConfig clientConfig, ServerIntrospector serverIntrospector,
                                      LoadBalancedBackOffPolicyFactory loadBalancedBackOffPolicyFactory,
                                      LoadBalancedRetryListenerFactory loadBalancedRetryListenerFactory, SpringClientFactory clientFactory) {
        super(lb, clientConfig, serverIntrospector);
        this.setRetryHandler(new DefaultLoadBalancerRetryHandler(clientConfig));
        this.loadBalancedBackOffPolicyFactory = loadBalancedBackOffPolicyFactory == null ? new LoadBalancedBackOffPolicyFactory.NoBackOffPolicyFactory() : loadBalancedBackOffPolicyFactory;
        this.loadBalancedRetryListenerFactory = loadBalancedRetryListenerFactory == null ? new LoadBalancedRetryListenerFactory.DefaultRetryListenerFactory() : loadBalancedRetryListenerFactory;
        this.clientFactory = clientFactory;
    }

    @Override
    public RibbonResponse execute(final RibbonRequest request, IClientConfig configOverride)
            throws IOException {
        final Request.Options options;
        if (configOverride != null) {
            options = new Request.Options(configOverride.get(CommonClientConfigKey.ConnectTimeout, this.connectTimeout),
                    (configOverride.get(CommonClientConfigKey.ReadTimeout, this.readTimeout)));
        } else {
            options = new Request.Options(this.connectTimeout, this.readTimeout);
        }
        final LoadBalancedRetryPolicy retryPolicy = new RibbonLoadBalancedRetryPolicy(this.getClientName(), this.clientFactory.getLoadBalancerContext(this.getClientName()), this);
        RetryTemplate retryTemplate = new RetryTemplate();
        BackOffPolicy backOffPolicy = loadBalancedBackOffPolicyFactory.createBackOffPolicy(this.getClientName());
        retryTemplate.setBackOffPolicy(backOffPolicy == null ? new NoBackOffPolicy() : backOffPolicy);
        RetryListener[] retryListeners = this.loadBalancedRetryListenerFactory.createRetryListeners(this.getClientName());
        if (retryListeners != null && retryListeners.length != 0) {
            retryTemplate.setListeners(retryListeners);
        }
        retryTemplate.setRetryPolicy(new FeignRetryPolicy(request.toHttpRequest(), retryPolicy, this, this.getClientName()));
        return retryTemplate.execute(retryContext -> {
            Request feignRequest = null;
            if (retryContext instanceof LoadBalancedRetryContext) {
                ServiceInstance service = ((LoadBalancedRetryContext) retryContext).getServiceInstance();
                if (service != null) {
                    log.info("feignRibbon targetServer={}:{}, retryCount:{}", service.getHost(), service.getPort(), retryContext.getRetryCount());
                    feignRequest = ((RibbonRequest) request.replaceUri(reconstructURIWithServer(new Server(service.getHost(), service.getPort()), request.getUri()))).toRequest();
                }
            }
            if (feignRequest == null) {
                feignRequest = request.toRequest();
            }
            Response response = request.client().execute(feignRequest, options);
            if (retryPolicy.retryableStatusCode(response.status())) {
                byte[] byteArray = response.body() == null ? new byte[]{} : StreamUtils.copyToByteArray(response.body().asInputStream());
                response.close();
                throw new RibbonResponseStatusCodeException(RetryableFeignLoadBalancer.this.clientName, response, byteArray, request.getUri());
            }
            return new RibbonResponse(request.getUri(), response);
        }, new RibbonRecoveryCallback<RibbonResponse, Response>() {
            @Override
            protected RibbonResponse createResponse(Response response, URI uri) {
                return new RibbonResponse(uri, response);
            }
        });
    }

    @Override
    public RequestSpecificRetryHandler getRequestSpecificRetryHandler(FeignLoadBalancer.RibbonRequest request, IClientConfig requestConfig) {
        return new RequestSpecificRetryHandler(false, false, this.getRetryHandler(), requestConfig);
    }

    @Override
    public ServiceInstance choose(String serviceId) {
        return new RibbonLoadBalancerClient.RibbonServer(serviceId, this.getLoadBalancer().chooseServer(serviceId));
    }
}